import numpy as np
import tensorflow as tf
# 引入TensorFlow的minit例子数据
from tensorflow.examples.tutorials.mnist import input_data
from matplotlib import pyplot as plt

tf.logging.set_verbosity(tf.logging.INFO)
# 首先查看数据内容:训练集、验证集和测试集
mnist = input_data.read_data_sets("./", one_hot=True)
# print(mnist.train.images.shape)
# print(mnist.train.labels.shape)
#
# print(mnist.validation.images.shape)
# print(mnist.validation.labels.shape)
#
# print(mnist.test.images.shape)
# print(mnist.test.labels.shape)
'''
结果：
(55000, 784)
(55000, 10)
(5000, 784)
(5000, 10)
(10000, 784)
(10000, 10)
'''
# 经过独热编码后，每张图片都是长度为28*28的一维向量，所以在用的时候需要将它们先还原成28*28的二维图片
plt.figure(figsize=(8, 8))
for idx in range(16):
    plt.subplot(4, 4, idx + 1)
    plt.axis('off')
    plt.title('[{}]'.format(np.argmax(mnist.train.labels[idx])))
    plt.imshow(mnist.train.images[idx].reshape((28, 28)))
plt.show()


'''
接下来，定义用于训练的网络，首先定义网络的输入。
这里我们直接使用上面的数据作为输入，所以定义两个placeholder分别用于图像和lable数据，
另外，定义一个bool类型的变量用于标识当前网络是否正在训练。为了让网络更高效的运行，
多个数据会被组织成一个batch送入网络，两个placeholder的第一个维度就是batchsize，因为
我们这里还没有确定batchsize，所以第一个维度留空。
'''
x = tf.placeholder("float", [None, 784], name='x')
y = tf.placeholder("float", [None, 10], name='y')
'''
因为我们输入的是图片展开后的一维向量，所以第一步就需要先把一维向量还原为二维的图片。
'''
x_image = tf.reshape(x, [-1, 28, 28, 1])
'''
接下来，我们定义第一个卷积层，使用6个5X5的卷积核对输入数据进行卷积， padding方式选择
valid，所以输出数据的宽高变为24x24,但是深度已经从原来的1变成了6。 本层卷积的激活函数
为relu。
'''
with tf.name_scope('conv1'):
    C1 = tf.contrib.slim.conv2d(x_image, 6, [5, 5], padding='VALID',
                                activation_fn=tf.nn.relu)
'''
接下来进行stride为2的最大池化，池化后，输出深度不变，但是长宽减半，所以输出变成了12x12,
深度6.
'''
with tf.name_scope('pool1'):
    S2 = tf.contrib.slim.max_pool2d(C1, [2, 2], stride=[2, 2], padding='VALID')
'''
接下来，我们定义第二个卷积层，使用16个5X5的卷积核对输入数据进行卷积， padding方式还是选
择valid，输出8x8,深度为16，本层卷积的激活函数为relu。
'''
with tf.name_scope('conv2'):
    C3 = tf.contrib.slim.conv2d(S2, 16, [5, 5], padding='VALID', activation_fn=tf.nn.relu)

'''
再进行一次stride为2的最大池化，输出为4x4,深度16。
'''
with tf.name_scope('pool2'):
    S4 = tf.contrib.slim.max_pool2d(C3, [2, 2], stride=[2, 2], padding='VALID')
'''
池化后的数据是3维的，这里做一个拉平的操作，将3维数据展开到1维，然后送入两层全连接，全连
接隐层中神经元个数分别为120，84。
'''
with tf.name_scope('fc1'):
    S4_flat = tf.contrib.slim.flatten(S4)
    C5 = tf.contrib.slim.fully_connected(S4_flat, 120, activation_fn=tf.nn.relu)
with tf.name_scope('fc2'):
    F6 = tf.contrib.slim.fully_connected(C5, 84, activation_fn=tf.nn.relu)
'''
对特征添加一个0.6的dropout，以40%的概率丢弃特征中的某些数据，这样可以提高网络的推广能力，
减少过拟合的可能性。需要注意的是，dropout仅在训练的时候使用，验证的时候，需要关闭dropout，
所以验证时候的keep_prob是1.0。dropout的输出最终送入一个隐层为10的全连接层，这个全连接层
即为最后的分类器。
'''
with tf.name_scope('dropout'):
    keep_prob = tf.placeholder(name='keep_prob', dtype=tf.float32)
    F6_drop = tf.nn.dropout(F6, keep_prob)
with tf.name_scope('fc3'):
    logits = tf.contrib.slim.fully_connected(F6_drop, 10)
'''
接下来定义loss和用于优化网络的优化器。loss计算使用了
sparse_softmax_cross_entropy_with_logits, 这样做的好处是labels可以不用手动做one_hot省了
一些麻烦。这里使用了sgd优化器，学习率为0.3。
???试试看，增大减小学习率，换个优化器再进行训练会发生什么。???
'''
cross_entropy_loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits
                                    (logits=logits, labels=y))
l2_loss = tf.add_n([tf.nn.l2_loss(w)
                    for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)])
for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
    print(w.name)
    tf.summary.histogram(w.name, w)
total_loss = cross_entropy_loss + 7e-5*l2_loss
tf.summary.scalar('cross_entropy_loss',cross_entropy_loss)
tf.summary.scalar('l2_loss', l2_loss)
tf.summary.scalar('total_loss', total_loss)
optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.3).minimize(total_loss)
'''
需要注意的是，上面的网络，最后输出的是未经softmax的原始logits，而不是概率分布， 要想
看到概率分布，还需要做一下softmax。将输出的结果与正确结果进行对比，即可得到我们的网络
输出结果的准确率。
'''
pred = tf.nn.softmax(logits)
correct_pred = tf.equal(tf.argmax(y, 1), tf.argmax(logits, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
'''
saver用于保存或恢复训练的模型。
'''
batch_size = 100
training_step = 1100
saver = tf.train.Saver()
'''
以上定义的所有操作，均为计算图，也就是仅仅是定义了网络的结构，实际需要运行的话，还需要
创建一个session，并将数据填入网络中。
'''
merged = tf.summary.merge_all()
with tf.Session() as sess:
    writer = tf.summary.FileWriter("logs/", sess.graph)
    sess.run(tf.global_variables_initializer())
    #定义验证集和测试集
    validate_data = {x:mnist.validation.images,
                     y:mnist.validation.labels,
                     keep_prob:1.0}
    test_data = {x:mnist.test.images,
                 y:mnist.test.labels,
                 keep_prob:1.0}
    for i in range(training_step):
        xs, ys = mnist.train.next_batch(batch_size)
        _, loss, rs = sess.run(
            [optimizer, cross_entropy_loss, merged],
            feed_dict={
                x: xs,
                y: ys,
                keep_prob: 0.6
            })
        # _, loss, rs = sess.run([optimizer, cross_entropy_loss, merged],feed_dict={x:xs, y:ys, keep_prob:0.6})
        writer.add_summary(rs, i)
        # 每100次训练打印一次损失值与验证准确率
        if i>0 and i%100 == 0:
            validate_accurac = sess.run(accuracy, feed_dict=validate_data)
            print('after %d training steps, the loss is %g, the validation '
                  'accuracy is %g'%(i, loss, validate_accurac))
            saver.save(sess, './model.ckpt', global_step=i)
            print('the training is finish')
            # 最终的测试准确率
            acc = sess.run(accuracy, feed_dict=test_data)
            print('the test accuracy is:', acc)

with tf.Session() as sess:
    ckpt = tf.train.get_checkpoint_state('/')
    if ckpt and ckpt.model_checkpoint_path:
        saver.restore(sess, ckpt.model_checkpoint_path)
        final_pred, acc = sess.run([pred, accuracy], feed_dict={x:mnist.test.images,
                                                                y:mnist.test.labels,
                                                                keep_prob:1.0})
        orders = np.argsort(final_pred)
        plt.figure(figsize=(8, 8))
        print(acc)
        for idx in range(16):
            order = orders[idx, :][-1]
            prob = final_pred[idx, :][order]
            plt.subplot(4, 4, idx +1)
            plt.axis('off')
            plt.title('{}:[{}]-[{:.1f}%]'.format(np.argmax(mnist.test.labels[idx]),
                                                 order, prob*100))
            plt.show(mnist.test.images[idx].reshape((28, 28)))
        plt.show()
    else:
        pass


